# 适应度函数
import numpy as np

from src.ga.impl.output import get_output


def fitness(x, operator, pattern='MRE', tested='tf'):
    # 得到tf和pytorch输出的numpy数组
    tf_output, torch_output = get_output(x, operator)
    benchmark = (tf_output + torch_output) / 2
    if tested == 'tf':
        diff_matrix = tf_output - benchmark
        single_item_matrix = diff_matrix / tf_output
        # 取绝对值
        single_item_matrix_abs = np.maximum(single_item_matrix, -single_item_matrix)
    elif tested == 'torch':
        diff_matrix = torch_output - benchmark
        single_item_matrix = diff_matrix / torch_output
        # 取绝对值
        single_item_matrix_abs = np.maximum(single_item_matrix, -single_item_matrix)
    # 算MRE或MARE
    if pattern == 'MRE':
        return np.sum(single_item_matrix_abs) / np.size(single_item_matrix_abs)
    elif pattern == 'MARE':
        return np.max(single_item_matrix_abs)